package com.self.av;

import android.app.Activity;
import android.content.ContentResolver;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.ImageFormat;
import android.graphics.Rect;
import android.graphics.YuvImage;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaExtractor;
import android.media.MediaFormat;
import android.media.MediaMetadataRetriever;
import android.net.Uri;
import android.os.Bundle;
import android.os.Environment;
import android.util.Log;
import android.view.View;
import android.widget.Button;
import android.widget.ProgressBar;
import android.widget.TextView;
import android.widget.Toast;

import androidx.appcompat.app.AppCompatActivity;
import androidx.core.content.ContextCompat;

import com.google.mediapipe.framework.image.BitmapImageBuilder;
import com.google.mediapipe.framework.image.MPImage;
import com.google.mediapipe.tasks.components.containers.NormalizedLandmark;
import com.google.mediapipe.tasks.core.BaseOptions;
import com.google.mediapipe.tasks.core.Delegate;
import com.google.mediapipe.tasks.vision.core.RunningMode;
import com.google.mediapipe.tasks.vision.poselandmarker.PoseLandmarker;
import com.google.mediapipe.tasks.vision.poselandmarker.PoseLandmarkerResult;

import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;

public class PreprocessActivity extends AppCompatActivity {
    private String videoMd5;
    private Uri videoUri;
//    private PoseLandmarker poseLandmarker;

    private ProgressBar progressBar;
    private TextView fileNameTextView;
    private TextView statusTextView;
    private TextView progressTextView;
    private ExecutorService processingExecutor; // 添加后台执行器

    private Button cancelButton;
    private Button startButton;
    private boolean isCanceled = false;
    private boolean setColorFormat = false;
    private int totalFrames = 0;
    private AtomicInteger finishFrameCount = new AtomicInteger(0);
    private ConcurrentHashMap<Long, PoseLandmarker> poseLandmarkerMap = new ConcurrentHashMap<>();
    private HashMap<Integer, List<NormalizedLandmark>> preprocessData = new HashMap<>();

    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        // 设置透明主题
        setContentView(R.layout.activity_preprocess);
        getWindow().addFlags(android.view.WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON);

//        loadMPModel();

        progressBar = findViewById(R.id.progressBar);
        fileNameTextView = findViewById(R.id.fileNameTextView);
        statusTextView = findViewById(R.id.statusTextView);
        progressTextView = findViewById(R.id.progressTextView);
        cancelButton = findViewById(R.id.cancelButton);
        startButton = findViewById(R.id.startButton);

        // 获取传递的文件名
        String videoUriStr = getIntent().getStringExtra("videoUri");
        long durationMs = getIntent().getLongExtra("totalDurationMs", 0); // 总时长（毫秒）
        float frameRate = getIntent().getFloatExtra("frameRate", 0f); // 帧率（FPS）
        Log.e("@@@@@@", "durationMs:" + durationMs);
        Log.e("@@@@@@", "frameRate:" + frameRate);
        videoUri = Uri.parse(videoUriStr);
        videoMd5 = Utils.calculateMD5FromUri(getApplicationContext(), videoUri);
        ContentResolver contentResolver = getContentResolver();
        String fileName = Utils.getFileName(contentResolver, videoUri);
        if (videoUriStr != null) {
            statusTextView.setText("待处理");
            fileNameTextView.setText(fileName);
        }

        MediaMetadataRetriever retriever = new MediaMetadataRetriever();
        retriever.setDataSource(this, videoUri);

        // 获取视频基本信息
        totalFrames = (int)(durationMs * frameRate / 1000); // 总帧数
        long frameIntervalUs = (long)(1000000 / frameRate); // 帧间隔（微秒）

        runOnUiThread(() -> {
            progressBar.setMax(totalFrames);
            progressTextView.setText(String.format("0/%d 帧", totalFrames));
        });

        File saveDir = new File(getExternalFilesDir(Environment.DIRECTORY_DOCUMENTS), "av_pose_results");
        if (!saveDir.exists() && !saveDir.mkdirs()) {
            Log.e("SaveResult", "无法创建保存目录");
            runOnUiThread(() -> Toast.makeText(this, "无法创建保存目录", Toast.LENGTH_SHORT).show());
            return;
        }
        File resultFile = new File(saveDir, videoMd5 + ".json");

        // 使用后台线程避免阻塞UI（文件IO操作）
//        Executors.newSingleThreadExecutor().execute(() -> {
            if (resultFile.exists()) {
                preprocessData = Utils.loadPreprocessData(getApplicationContext(), videoUri);
                if (preprocessData.size() > totalFrames){
                    preprocessData.clear();
                }
                if (preprocessData.size() == totalFrames) {
                    statusTextView.setText("已完成，可重新预处理");
                } else if (preprocessData.size() < totalFrames && preprocessData.size() > 0) {
                    statusTextView.setText("未完成，可继续预处理");
                }
                finishFrameCount.addAndGet(preprocessData.size());
                int progress = finishFrameCount.get();
                runOnUiThread(() -> {
                    progressBar.setProgress(progress);
                    progressTextView.setText(String.format("%d/%d 帧", progress, totalFrames));
                });
            }
//        });

        startButton.setOnClickListener(v -> {
            // 新增：开始视频帧处理
            if (videoUriStr != null) {
                if (preprocessData.size() == totalFrames){
                    preprocessData.clear();
                    finishFrameCount.set(0);
                }
                startButton.setEnabled(false);
                startButton.setBackgroundColor(ContextCompat.getColor(getApplicationContext(), R.color.gray_125));
                statusTextView.setText("处理中");
                int threadCount = Runtime.getRuntime().availableProcessors();
                processingExecutor = Executors.newFixedThreadPool(threadCount);
                processingExecutor.execute(() -> processVideoFrames(retriever, totalFrames, frameIntervalUs, resultFile));
            } else {
                runOnUiThread(() -> Toast.makeText(this, "视频Uri为空", Toast.LENGTH_SHORT).show());
                finish();
            }
        });

        cancelButton.setOnClickListener(v -> {
            runOnUiThread(() -> Toast.makeText(this, "防止误触，长按取消退出", Toast.LENGTH_SHORT).show());
        });
        // 新增：长按取消处理
        cancelButton.setOnLongClickListener(v -> {
            isCanceled = true;
            // 新增：取消视频帧处理
            if (processingExecutor != null && !processingExecutor.isShutdown()) {
                processingExecutor.shutdownNow();
            }
            finish();
            return false;
        });
    }

    // 新增：视频帧处理核心方法
    private void processVideoFrames(MediaMetadataRetriever retriever, int totalFrames, long frameIntervalUs, File resultFile) {
        MediaExtractor extractor = null;
        MediaCodec codec = null;
        try {
//            int targetWidth = 640;  // 目标宽度，可根据需求调整
//            int targetHeight = 480; // 目标高度，可根据需求调整

            File saveDir = new File(getExternalFilesDir(Environment.DIRECTORY_DOCUMENTS), "av_pose_results");
            if (!saveDir.exists() && !saveDir.mkdirs()) {
                Log.e("SaveResult", "无法创建保存目录");
                runOnUiThread(() -> Toast.makeText(this, "无法创建保存目录", Toast.LENGTH_SHORT).show());
                return;
            }

            boolean isAppend = preprocessData.size() > 0;
            BufferedWriter writer = new BufferedWriter(new FileWriter(resultFile, isAppend));

            // 初始化 MediaExtractor 和 MediaCodec
            extractor = new MediaExtractor();
            extractor.setDataSource(getApplicationContext(), videoUri, null);

            // 查找视频轨道
            int videoTrackIndex = -1;
            for (int i = 0; i < extractor.getTrackCount(); i++) {
                MediaFormat format = extractor.getTrackFormat(i);
                String mime = format.getString(MediaFormat.KEY_MIME);
                if (mime.startsWith("video/")) {
                    videoTrackIndex = i;
                    break;
                }
            }
            if (videoTrackIndex == -1) {
                Log.e("MediaCodec", "未找到视频轨道");
                return;
            }

            // 配置 MediaCodec
            MediaFormat videoFormat = extractor.getTrackFormat(videoTrackIndex);
            extractor.selectTrack(videoTrackIndex);
            String mime = videoFormat.getString(MediaFormat.KEY_MIME);
            codec = MediaCodec.createDecoderByType(mime);
            codec.configure(videoFormat, null, null, 0);
            codec.start();

            MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
            boolean isInputDone = false;
            boolean isOutputDone = false;
            int frameCount = 0;

            // 新增：批次处理相关变量
            List<FrameData> frameBatch = new ArrayList<>();
            final int BATCH_SIZE = 100; // 每批处理100帧
            List<Future<?>> taskFutures = new ArrayList<>();

            // 循环处理 codec 输入/输出缓冲区
            while (!isOutputDone && !isCanceled && frameCount < totalFrames) {

                // 处理输入缓冲区（喂数据给解码器）
                if (!isInputDone) {
                    int inputBufferId = codec.dequeueInputBuffer(10000);
                    if (inputBufferId >= 0) {
                        ByteBuffer inputBuffer = codec.getInputBuffer(inputBufferId);
                        int sampleSize = extractor.readSampleData(inputBuffer, 0);
                        if (sampleSize < 0) {
                            codec.queueInputBuffer(inputBufferId, 0, 0, 0, MediaCodec.BUFFER_FLAG_END_OF_STREAM);
                            isInputDone = true;
                        } else {
                            long presentationTimeUs = extractor.getSampleTime();
                            codec.queueInputBuffer(inputBufferId, 0, sampleSize, presentationTimeUs, 0);
                            extractor.advance();
                        }
                    }
                }

                // 处理输出缓冲区（获取解码后的帧）
                int outputBufferId = codec.dequeueOutputBuffer(bufferInfo, 10000);
                if (outputBufferId >= 0) {
                    if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
                        isOutputDone = true;
                    }

                    // 获取解码后的帧数据并转换为 Bitmap
                    MediaFormat outputFormat = codec.getOutputFormat(outputBufferId);
                    ByteBuffer outputBuffer = codec.getOutputBuffer(outputBufferId);
                    Bitmap frameBitmap = decodeToBitmap(outputBuffer, outputFormat);
                    if (frameBitmap != null) {
                        // 新增：将帧添加到批次
                        if (!preprocessData.containsKey(frameCount)) {
                            frameBatch.add(new FrameData(frameCount, bufferInfo.presentationTimeUs, frameBitmap));
                        }

                        // 新增：当批次满100帧或处理完成时，提交批次任务
                        if (frameBatch.size() >= BATCH_SIZE || isOutputDone || frameCount >= totalFrames - 1) {
                            List<FrameData> batch = new ArrayList<>(frameBatch);
                            frameBatch.clear();

                            // 提交批次处理任务
                            Future<?> future = processingExecutor.submit(() -> processFrameBatch(batch, writer));
                            taskFutures.add(future);
                        }

                        // 更新进度（确保 frameCount 从 0 开始，与总帧数匹配）
                        frameCount++;

//                        frameBitmap.recycle();

                    }

                    codec.releaseOutputBuffer(outputBufferId, false);
                }
            }

            // 新增：等待所有批次任务完成
            for (Future<?> future : taskFutures) {
                if (isCanceled) {
                    break;
                }
                try {
                    future.get();
                } catch (InterruptedException | ExecutionException e) {
                    Log.e("ProcessError", "批次任务执行异常", e);
                }
            }

            // 处理完成
            runOnUiThread(() -> {
//                Toast.makeText(this, "视频预处理完成", Toast.LENGTH_SHORT).show();
//                setResult(Activity.RESULT_OK);
                finish();
            });

            retriever.release();
        } catch (Exception e) {
            Log.e("PreprocessError", "处理失败", e);
            runOnUiThread(() -> Toast.makeText(this, "处理失败: " + e.getMessage(), Toast.LENGTH_LONG).show());
        } finally {
            if (processingExecutor != null) {
                processingExecutor.shutdown();
            }
            for (PoseLandmarker landmarker : poseLandmarkerMap.values()) {
                if (landmarker != null) {
                    landmarker.close();
                }
            }
            poseLandmarkerMap.clear();
        }
    }


    // 新增：将 MediaCodec 输出的 YUV 帧转换为 Bitmap
    private Bitmap decodeToBitmap(ByteBuffer buffer, MediaFormat format) {
        int width = format.getInteger(MediaFormat.KEY_WIDTH);
        int height = format.getInteger(MediaFormat.KEY_HEIGHT);
        int targetWidth = width / 2;
        int targetHeight = height / 2;
        int colorFormat = format.getInteger(MediaFormat.KEY_COLOR_FORMAT);
        if (!setColorFormat) {
            runOnUiThread(() -> Toast.makeText(this, "色彩格式: " + colorFormat, Toast.LENGTH_LONG).show());
            setColorFormat = true;
        }

        // 仅处理 YUV420 系列格式（包含完整色彩信息）
        if (colorFormat == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible ||
                colorFormat == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar ||
                colorFormat == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar) {

            buffer.rewind();
            byte[] yuvData = new byte[buffer.remaining()];
            buffer.get(yuvData);

            int ySize = width * height;
            int uvSize = ySize / 4; // U/V 分量总大小（各占 ySize/4）

            // 1. 处理 YUV420Planar (I420 格式：Y平面 + U平面 + V平面) → 转换为 NV21
            if (colorFormat == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar) {
                byte[] nv21Data = new byte[ySize + uvSize * 2]; // NV21 总大小 = Y + VU交织平面
                // 复制 Y 平面（保持不变）
                System.arraycopy(yuvData, 0, nv21Data, 0, ySize);
                // 提取 U/V 平面数据（I420 中 U 在 V 前）
                byte[] uData = new byte[uvSize];
                byte[] vData = new byte[uvSize];
                System.arraycopy(yuvData, ySize, uData, 0, uvSize); // U 平面（ySize ~ ySize+uvSize）
                System.arraycopy(yuvData, ySize + uvSize, vData, 0, uvSize); // V 平面（ySize+uvSize ~ 末尾）
                // 交织 V/U 分量（NV21 要求 VU 交织）
                int nv21Index = ySize;
                for (int i = 0; i < uvSize; i++) {
                    nv21Data[nv21Index++] = vData[i]; // V 分量在前
                    nv21Data[nv21Index++] = uData[i]; // U 分量在后
                }
                yuvData = nv21Data; // 更新为 NV21 格式数据
            }
            // 2. 处理 YUV420SemiPlanar (NV12 格式：Y平面 + UV交织平面) → 转换为 NV21
            else if (colorFormat == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar) {
                // 交换 UV 分量（NV12: Y-UV交织 → NV21: Y-VU交织）
                for (int i = ySize; i < yuvData.length; i += 2) {
                    byte temp = yuvData[i];
                    yuvData[i] = yuvData[i + 1];
                    yuvData[i + 1] = temp;
                }
            }

            // 配置 NV21 格式参数（YuvImage 仅支持 NV21 和 YUY2）
            int[] strides = new int[]{width, width}; // Y平面每行宽度、VU平面每行宽度
            YuvImage yuvImage = new YuvImage(yuvData, ImageFormat.NV21, width, height, strides);

            // 压缩为 JPEG 并解码为 RGB Bitmap
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            yuvImage.compressToJpeg(new Rect(0, 0, width, height), 70, outputStream);
            byte[] rgbData = outputStream.toByteArray();
            Bitmap rgbBitmap = BitmapFactory.decodeByteArray(rgbData, 0, rgbData.length);
//            return rgbBitmap;
            // 缩放到目标分辨率
            return Bitmap.createScaledBitmap(rgbBitmap, targetWidth, targetHeight, true);
        }
        return null;
    }


    public PoseLandmarker loadMPModel() {
        // 模型文件在assets目录下的文件名
        long threadId = Thread.currentThread().getId();
        PoseLandmarker landmarker = poseLandmarkerMap.get(threadId);
        if (landmarker == null) {
            Log.e("initMPModel", "初始化MP模型: " + threadId);
            BaseOptions.Builder baseOptionsBuilder = BaseOptions.builder();
            baseOptionsBuilder.setDelegate(Delegate.GPU);
            baseOptionsBuilder.setModelAssetPath("pose_landmarker_heavy.task");
            PoseLandmarker.PoseLandmarkerOptions.Builder optionsBuilder = PoseLandmarker.PoseLandmarkerOptions.builder();
            optionsBuilder.setBaseOptions(baseOptionsBuilder.build())
                    .setMinPoseDetectionConfidence(0.8f)
                    .setMinTrackingConfidence(0.9f)
                    .setMinPosePresenceConfidence(0.8f)
                    .setNumPoses(1)
                    .setOutputSegmentationMasks(false)
                    .setRunningMode(RunningMode.VIDEO);
            PoseLandmarker.PoseLandmarkerOptions poseLandmarkerOptions = optionsBuilder.build();
            landmarker = PoseLandmarker.createFromOptions(getApplicationContext(), poseLandmarkerOptions);
            poseLandmarkerMap.put(threadId, landmarker);
        } else {
            Log.e("initMPModel", "MP模型已存在: " + threadId);
        }
        return landmarker;
    }


    // 新增：帧数据封装类
    private static class FrameData {
        int frameIndex;
        long timestampUs;
        Bitmap frameBitmap;

        FrameData(int frameIndex, long timestampUs, Bitmap frameBitmap) {
            this.frameIndex = frameIndex;
            this.timestampUs = timestampUs;
            this.frameBitmap = frameBitmap;
        }
    }

    // 新增：批次处理方法
    private void processFrameBatch(List<FrameData> batch, BufferedWriter writer) {
        if (isCanceled) return;
        PoseLandmarker poseLandmarker = loadMPModel();
        for (FrameData frameData : batch) {
            try {
                if (isCanceled) return;
                // 处理当前帧
                MPImage mpImage = new BitmapImageBuilder(frameData.frameBitmap).build();
                PoseLandmarkerResult result = poseLandmarker.detectForVideo(mpImage, frameData.timestampUs / 1000);

                JSONObject frameJson = new JSONObject();
                if (result.landmarks().size() > 0) {

                    frameJson.put("frameIndex", frameData.frameIndex);
                    frameJson.put("timestampUs", frameData.timestampUs);

                    JSONArray landmarksArray = new JSONArray();
                    for (List<NormalizedLandmark> landmarks : result.landmarks()) {
                        JSONArray poseArray = new JSONArray();
                        for (NormalizedLandmark landmark : landmarks) {
                            JSONObject landmarkJson = new JSONObject();
                            landmarkJson.put("x", landmark.x());
                            landmarkJson.put("y", landmark.y());
                            landmarkJson.put("z", landmark.z());
                            landmarkJson.put("visibility", landmark.visibility().get());
                            poseArray.put(landmarkJson);
                        }
                        landmarksArray.put(poseArray);
                    }
                    frameJson.put("landmarks", landmarksArray);

                } else {
                    frameJson.put("frameIndex", frameData.frameIndex);
                    frameJson.put("timestampUs", frameData.timestampUs);
                    frameJson.put("landmarks", new JSONArray());
//                    Utils.save_local(getApplicationContext(), frameData.frameBitmap);
                }
                // 同步写入结果，确保线程安全
                synchronized (writer) {
                    writer.write(frameJson.toString() + "\r\n");
                    writer.flush();
                }
                frameData.frameBitmap.recycle();

                int progress = finishFrameCount.incrementAndGet();
                runOnUiThread(() -> {
                    progressBar.setProgress(progress);
                    progressTextView.setText(String.format("%d/%d 帧", progress, totalFrames));
                });
            } catch (JSONException | IOException e) {
                Log.e("SaveResultError", "写入帧结果失败", e);
            } finally {
                frameData.frameBitmap.recycle();
            }
        }
//        poseLandmarker.close();
    }


    @Override
    public void onBackPressed() {
        // 执行与“取消”按钮相同的清理逻辑
//        isCanceled = true;
//        if (processingExecutor != null && !processingExecutor.isShutdown()) {
//            processingExecutor.shutdownNow(); // 停止后台任务
//        }
//        finish(); // 销毁当前Activity
        runOnUiThread(() -> Toast.makeText(this, "防止误触，长按取消退出", Toast.LENGTH_SHORT).show());
    }
    @Override
    protected void onDestroy() {
        super.onDestroy();
//        if (poseLandmarker != null) {
//            poseLandmarker.close();
//        }
        if (processingExecutor != null && !processingExecutor.isShutdown()) {
            processingExecutor.shutdownNow();
        }
    }

}